#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>

Last Changed on: $Date$ 
'''

__revision__ = '$Revision$'

import sys
try:
    import psyco
    psyco.full()
except ImportError:
    pass

from sympy import Poly, sympify, Symbol, diff
#Non-standard modules
from mrange import cumRange
from VariableDict import Variables
from GetMomentEquation import getMomentEquation, getCumulantEquation
import Maple
from GetReactionInformation import GetStoichiometryDict
from GlobalVariables import globalVariables
from GetEvents import GetEvents
from GetTimeDerivatives import getTimeDerivative


class MapleInformation(object):
    '''Used as a storage container'''
    
    def __init__(self):
        self.pars = []
        self.species = []
        self.mom_equs = []
        self.time_derivatives = []

        
    def putFileName(self, file_name):
        self.file_name = file_name
        self.fp = open(self.file_name, 'w')
    
    def getFile(self):
        return self.fp

    def closeFile(self):
        self.fp.close()
        return    
    
    def putListOfSpecies(self, list_of_species):
        self.species = list_of_species
    
    def getListOfSpecies(self):
        return self.species
    
    def putStoicDict(self, gsd):
        self.gsd = gsd
     
    def getStoicDict(self):
        return self.gsd

    def putListOfPars(self, list_of_pars):
        self.pars = list_of_pars

    def getListOfPars(self):
        return self.pars
    
    def putListOfMomEqus(self, list_of_mom_equs):
        self.mom_equs = list_of_mom_equs
    
    def getListOfMomEqus(self):
        return self.mom_equs
    

def writeHeader(maple_information):
    '''Writes the maple file header.'''
    m = maple_information
    
    m.getFile().write(Maple.getHeader())
    expression = Maple.wrapExpression('restart;')
    m.getFile().write(expression)
    expression = Maple.wrapExpression('infolevel[dsolve]:=2;')
    m.getFile().write(expression)
    return True

def writeParameters(maple_information):
    m = maple_information
    
    for pa in m.getStoicDict().pa_dict.keys():
       # print '...Getting the parameter %s' % pa
        m.pars.append(pa + ':=' + str(m.getStoicDict().pa_dict[pa]['value']))
    

    #Get ready for outputing into maple   
    section_header = Maple.getSectionHeader('Constant Species and Parameters')
    m.getFile().write(section_header)
    maple_pars = ";".join(m.pars) + ";"    
    expression = Maple.wrapExpression(maple_pars)
    m.getFile().write(expression)
    section_footer = Maple.getSectionFooter()
    m.getFile().write(section_footer)
    return True

def writeDifferentialEquations(maple_information):
    m = maple_information
    
    section_header = Maple.getSectionHeader('ODEs')
    m.getFile().write(section_header)
    i = 0
    eq_list = []
    expression = ''
    for species, reaction in zip(m.species, m.mom_equs):
        expression += 'eq_%s:= diff(%s, t)= %s;\n' % (i, str(Variables().getMoment(species)), reaction)
        eq_list.append('eq_' + str(i))
        i += 1

    expression = Maple.wrapExpression(expression)
    m.getFile().write(expression)

    section_footer = Maple.getSectionFooter()
    m.getFile().write(section_footer)
    return tuple(eq_list)


def getEqualsTo(m, specie_tup, eid):

    #eid == event_info_dict
    event_no = eid['event_no']
    time = eid['time']
    assignments = eid['assignments']
    assign_values = eid['assign_values']

    #First check for assignments
    sps_in_ass = {}
    for i, value in enumerate(specie_tup):
        if value != 0:
            sp = m.getStoicDict().m.getSpecies(i)
            if sp.getId() in assignments:
                sps_in_ass[i] = value

    specie_list = list(specie_tup)
    for idx in sps_in_ass.keys():
        specie_list[idx] = 0
    specie_tup = tuple(specie_list)

   
    equals_to = ''
    if sum(specie_tup) != 0:
        maple_index = m.getListOfSpecies().index(specie_tup) + 1
        equals_to = 'h%s(%s)[%s]' % (str(int(event_no)-1), time, maple_index)

    #Now add in the bits from the assignments
    for idx, power in sps_in_ass.items():
            sp = m.getStoicDict().m.getSpecies(idx)
            print eid
            ass_idx = assignments.index(sp.getId())
            sp_value = str(assign_values[ass_idx])
            
            if len(equals_to):
                equals_to += '*' + sp_value
            else:
                equals_to = sp_value
            if power > 1:
                equals_to += '**%s' % str(power)
    equals_to = '(%s) = %s' % (time, equals_to)
    return equals_to




def getEqualsTo4Tequal0(m, specie_tup):
    equals_to = ''
    for i, value in enumerate(specie_tup):
        if value != 0:
            sp = m.getStoicDict().m.getSpecies(i)
            sp_value = str(sp.getInitialAmount())
            if len(equals_to):
                equals_to += '*' + sp_value
            else:
                equals_to = sp_value
            if value > 1:
                equals_to += '**%s' % str(value)
    equals_to = '(0) = ' + equals_to
    return equals_to




def writeMapleSpeciesAtEvent(m, eq_tup, eid):
    #eid == event_info_dict
    event_no = eid['event_no']
    time = eid['time']
    assignments = eid['assignments']
    assign_values = eid['assign_values']
    
    maple_species = []
    eq_list = list(eq_tup)
   
    for i, species in enumerate(m.species):
        t_species = tuple(species)
       
        if sum(species) > 1: #higher order moment

            if time == '0':
                equals_to = getEqualsTo4Tequal0(m, species)
                maple_sp = str(Variables().getMoment(t_species))[:-3] + equals_to
            else:

                equals_to = getEqualsTo(m, species, eid)
                maple_sp = str(Variables().getMoment(t_species))[:-3] + equals_to   

                print equals_to, species
             
        else:
            if time == '0':
                equals_to = getEqualsTo4Tequal0(m, species)
                maple_sp = str(Variables().getMoment(t_species))[:-3] + equals_to
            else:
                equals_to = getEqualsTo(m, species, eid)
                maple_sp = str(Variables().getMoment(t_species))[:-3] + equals_to
           
           
        maple_species.append(str(Variables().getMoment(t_species)))
        eq_list.append(maple_sp)
    
    maple_initial_conditions = 'IC' + event_no + ':={' + ", ".join(eq_list) + '};'
    expression = Maple.wrapExpression(maple_initial_conditions)
    m.getFile().write(expression)

    maple_species = ",".join(maple_species)
    
    part =  'F' + event_no + ':=dsolve(IC' + event_no + ', {' + maple_species + '}, type=numeric, method=rosenbrock);'
    expression = Maple.wrapExpression(part)
    m.getFile().write(expression) 
    
    part = 'h' + event_no + ':= theta->eval([' + maple_species+ '],F' + event_no + '(theta));'
    expression = Maple.wrapExpression(part)
    m.getFile().write(expression)
     
    return True


def writeEvents(maple_information, eq_tup):
    m = maple_information
    
    #event_info_dict == eid
    eid = {'event_no':'0', 'time':'0', 'assignments':[], 'assign_values':[]}
    writeMapleSpeciesAtEvent(m, eq_tup, eid)


    events, order = GetEvents(m.getStoicDict().m).getEvents()
    no_of_events = len(events)
    for i, event in enumerate(order):
        eid['time'] = events[event]['trigger'][2]
        eid['assignments'] = events[event]['assignments']
        eid['assign_values'] = events[event]['values']
        eid['event_no'] = str(i+1)
        writeMapleSpeciesAtEvent(m, eq_tup, eid)

    

def createMapleFile(sbml_file, output_file='/tmp/tmp.mw'):
    '''Creates the maple file with moment equations.
    
    Arguments:
        sbml_file - sbml file name
        output_file - outout file
    '''
    
    gsd = GetStoichiometryDict(sbml_file = sbml_file)
    m = MapleInformation()
    m.putFileName(output_file)
    m.putStoicDict(gsd)
    
    reaction_list = gsd.getReactionList()
    max_values = [globalVariables()['closure_order']] * len(gsd.empty_list)

    for values in cumRange(max_values):
        print values
        mom_eq = getMomentEquation(values, reaction_list, gsd.hasDenominator)
        print 1
        time_derivative = getTimeDerivative(values, gsd)
        moment = Variables().getMoment(values)
        t_d = diff(moment, sympify('t'))
        print "#################"
        print time_derivative
        print mom_eq
        print "##################"
        
        #mom_eq = mom_eq - (time_derivative - t_d * time_derivative.coeff(t_d)).expand()
        mom_eq = mom_eq / time_derivative.coeff(t_d)
        #print values, globalVariables()['closure_order'], globalVariables()['time_closure_order']
        #print time_derivative
        #print mom_eq
        #if mom_eq != 0:
        m.species.append(values)
        #print '...Getting moment equation  dkappa_%s/dt' % str(values)
        m.mom_equs.append(str(mom_eq))

        
#        else:
 #           print '0:=', values
#            if sum(values) > 1: #XXX Is this if statement needed?
#                print 'Yes, this statement is need', values
#                equals_to = getEqualsTo(m, values)
#                maple_sp = str(Variables().getMoment(tuple(values)))[:-3] + equals_to
#                m.pars.append(maple_sp)
#            else:
#                index = list(values).index(1)
#                sp_id = gsd.index2species[index]
#                print 'Note: Species %s is constant.' % sp_id
#                print 'If you add - constant ="true" - to your SBML file then you will have fewer moment equations'
#                
#                m.species.append(values)
#                m.mom_equs.append(str(mom_eq))
                #sys.exit()
                
   

    writeHeader(m)
    writeParameters(m)
    eq_tup = writeDifferentialEquations(m)
    writeEvents(m, eq_tup) 
    
    m.getFile().write(Maple.getFooter())
    m.closeFile()
    return True



if __name__ == '__main__':
    globalVariables(output_format = 'Maple')
#    createMapleFile('../xml_files/imdeath.xml')
    createMapleFile('xml_files/mm.xml')    
 #   createMapleFile('/tmp/BIOMD0000000012.xml')
 

    #alpha_dict = GetLambdaDict(sbml_file='xml_files/sis.xml').GetLambdaDict()
#    print getDifferentialEquations(sbml_file='../xml_files/Hsp90model1.xml', moment_order = 2)
#    print getDifferentialEquations(sbml_file='xml_files/tmp.xml', moment_order = 2)
#    print getDifferentialEquations(sbml_file='xml_files/sis.xml', moment_order = 2)

    
    
    

    print 'EOF'
#h:=theta->eval(kappa_00000000000002000(t),F(theta));
 
#plot(h, 0..300);
